# -*- coding: utf-8 -*-
# @Author  : 张佳玉
# @Time    : 2023/7/10 9:31
# @Function: LSB


import matplotlib.pyplot as plt
import matplotlib.image as pmg
import numpy as np

# RGB转灰度图
def rgb2gray(rgb):
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    gray = 0.2989 * r + 0.5870 * g + 0.1140 * b

    return np.trunc(gray)  # np.trunc对矩阵取整


#读入图像
limit=200#二值化临界值
cr="copyright.jpg"#copyright.jpg
copy_right=pmg.imread(cr)
origin=pmg.imread("LenaRGB.bmp")

if copy_right.shape[0] > origin.shape[0] or \
                    copy_right.shape[1] > origin.shape[1]:
        raise ValueError("copyright image must be smaller than original image")
#显示读入的图像
plt.imshow(copy_right)
plt.title("copyright_picture")
plt.axis('off')
plt.figure()
plt.show()

plt.imshow(origin)
plt.axis('off')
plt.show()

plt.title("origin_picture")
#图片备份
watermark=origin.copy()
copyrig=copy_right.copy()
#对版权图片二值化
for m in range(3):
    for i in range(copyrig.shape[0]):
        for j in range(copyrig.shape[1]):
            if copyrig[i,j,m]>limit:
                copyrig[i,j,m]=1
            else:
               copyrig[i,j,m]=0
#对原图最小位置零
for i in range(0,watermark.shape[0]):
    for j in range(0,watermark.shape[1]):
        watermark[i,j,:]=(watermark[i,j,:]//2)*2   #//取整除运算符
#将二值化的版权图片加入原图中
h,w=copyrig.shape[0],copyrig.shape[1]
watermark[0:h,0:w,:]=watermark[0:h,0:w,:]+copyrig
#显示加水印后的图像
plt.figure()
plt.imshow(watermark)
plt.axis('off')
plt.title("watermark")
plt.show()

#提取加水印图像的信息
c_watermark=watermark.copy()
watermark_copyright=(c_watermark%2)*255
#显示提取出的水印信息
plt.figure()
plt.imshow(watermark_copyright)
plt.axis('off')
plt.title("watermark_copyright")
plt.show()